use std::{fs, io, path::Path};

use crate::util::{property, rerank};
use colored::Colorize;
use itertools::Itertools;
use rig::{completion::ToolDefinition, tool::Tool};
use rig::{
    providers::{self, ollama::Client},
    vector_store::VectorStoreIndexDyn,
    Embed, OneOrMany,
};
use rig_sqlite::{Column, ColumnValue, SqliteVectorStore, SqliteVectorStoreTable};
use rusqlite::ffi::sqlite3_auto_extension;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sqlite_vec::sqlite3_vec_init;
use tokio::task;
use tokio_rusqlite::Connection;
use tree_sitter::{Parser, TreeCursor};
use walkdir::WalkDir;

#[derive(Deserialize)]
pub struct GetByClassArgs {
    class_name: String,
    method_name: Option<String>,
}

#[derive(Debug, thiserror::Error)]
#[error("Math error")]
pub struct JavaParseError;

#[derive(Deserialize, Serialize)]
pub struct JavaParseCmd {
    pub project_path: String,
}

impl Tool for JavaParseCmd {
    const NAME: &'static str = "java_parser";

    type Error = JavaParseError;
    type Args = GetByClassArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "java_parser".to_string(),
            description: "Java项目解析，根据类名和方法名精确查询代码".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "class_name": {
                        "type": "string",
                        "description": "类的全路径名"
                    }
                }
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        let project_path = &self.project_path;
        let method_name = args.method_name.as_ref();
        let class_name = args.class_name;

        // 查询结果
        let result: Vec<String> =
            by_class_name_and_method(project_path.as_str(), class_name.as_str(), method_name);

        Ok(result.join("\n"))
    }
}

fn by_class_name_and_method(
    project_path: &str,
    class_name: &str,
    method_name: Option<&String>,
) -> Vec<String> {
    // java 解析器
    let mut parser = Parser::new();
    parser
        .set_language(&tree_sitter_java::LANGUAGE.into())
        .expect("tree_sitter_sequel init fail");

    let contents: Vec<String> = WalkDir::new(project_path)
        .sort_by_file_name()
        .follow_links(true)
        .into_iter()
        .flat_map(|entry| {
            if let Ok(path) = entry {
                let full_path: String = path.path().display().to_string();
                if full_path.ends_with(".java") {
                    if let Ok(file_content) = fs::read_to_string(&full_path) {
                        // 解析
                        let node_tree = parser.parse(file_content.as_str(), None).unwrap();

                        // 查找package
                        let mut node_cursor = node_tree.walk();
                        if let Some(package_name) =
                            get_package_name(&mut node_cursor, &file_content)
                        {
                            let package_with_dot = format!("{}.", package_name);
                            if class_name.starts_with(&package_with_dot) {
                                // 当前文件可能有指定类
                                // 截取类名
                                if let Some(class_name_short) =
                                    class_name.strip_prefix(&package_with_dot)
                                {
                                    if let Some(class_content) = get_class_by_name(
                                        &class_name_short,
                                        &mut node_cursor,
                                        &file_content,
                                    ) {
                                        // 查找制定方法
                                        if let Some(method_name) = &method_name {
                                            let method_content = get_method_by_name(
                                                method_name,
                                                &mut node_cursor,
                                                &file_content,
                                            );
                                            return method_content;
                                        } else {
                                            return vec![class_content];
                                        }
                                    }
                                }
                            }
                        }

                        return vec![];
                    }
                }
            }
            return vec![];
        })
        .collect();
    return contents;
}

fn get_method_by_name(
    method_name: &str,
    node_cursor: &mut TreeCursor,
    source_code: &str,
) -> Vec<String> {
    let root_node = node_cursor.node();
    let mut methods: Vec<String> = vec![];
    loop {
        let current_node = node_cursor.node();
        if current_node.kind() == "formal_parameters" {
            // 逆向寻找 方法名,返回值
            node_cursor.goto_previous_sibling();
            node_cursor.goto_previous_sibling();

            // 返回值
            let return_type_node = node_cursor.node();
            let return_type =
                &source_code[return_type_node.start_byte()..return_type_node.end_byte()];

            // 方法名
            node_cursor.goto_next_sibling();
            let method_name_node = node_cursor.node();
            let now_method_name =
                &source_code[method_name_node.start_byte()..method_name_node.end_byte()];

            if now_method_name == method_name {
                // 找到指定方法
                let mut method_content = vec![return_type, now_method_name];

                loop {
                    // 参数 异常 方法体 ;
                    node_cursor.goto_next_sibling();
                    let now_node = node_cursor.node();
                    let now_node_content = &source_code[now_node.start_byte()..now_node.end_byte()];
                    method_content.push(now_node_content);

                    if now_node.kind() == "block" || now_node.kind() == ";" {
                        break;
                    }
                }

                // 保存
                methods.push(method_content.join(" "));

                // 进入下一个节点
                node_cursor.goto_next_sibling();
            } else {
                // 方法名不匹配，回到 formal_parameters
                node_cursor.goto_next_sibling();
            }
        }

        if node_cursor.goto_first_child() {
            continue;
        }

        loop {
            if node_cursor.goto_next_sibling() {
                break;
            }

            node_cursor.goto_parent();
            if root_node == node_cursor.node() {
                return methods;
            }
        }
    }
}

fn get_class_by_name(
    class_name: &str,
    node_cursor: &mut TreeCursor,
    source_code: &str,
) -> Option<String> {
    loop {
        let current_node = node_cursor.node();
        if current_node.kind() == "class" {
            let class_class = &source_code[current_node.start_byte()..current_node.end_byte()];
            // 移动到兄弟节点 className
            if node_cursor.goto_next_sibling() {
                let class_name_node = node_cursor.node();
                if class_name_node.kind() == "identifier" {
                    let class_identifier =
                        &source_code[class_name_node.start_byte()..class_name_node.end_byte()];
                    let node_class_name =
                        &source_code[class_name_node.start_byte()..class_name_node.end_byte()];

                    if class_name == node_class_name {
                        // 找到指定class
                        if node_cursor.goto_next_sibling() {
                            let class_content_node = node_cursor.node();
                            let class_content = &source_code
                                [class_content_node.start_byte()..class_content_node.end_byte()];
                            return Some(format!(
                                "{class_class} {class_identifier} {class_content}"
                            ));
                        } else {
                            // 回到identifier
                            node_cursor.goto_previous_sibling();
                        }
                    }
                }
            }

            // 一个文件只有一个packgeName
            return None;
        }

        if node_cursor.goto_first_child() {
            continue;
        }

        loop {
            if node_cursor.goto_next_sibling() {
                break;
            }

            if !node_cursor.goto_parent() {
                return None;
            }
        }
    }
}

fn get_package_name(node_cursor: &mut TreeCursor, source_code: &str) -> Option<String> {
    loop {
        let current_node = node_cursor.node();
        if current_node.kind() == "package" {
            // 移动到兄弟节点
            if node_cursor.goto_next_sibling() {
                let scoped_identifier_node = node_cursor.node();
                if scoped_identifier_node.kind() == "scoped_identifier" {
                    let scoped_identifier = &source_code
                        [scoped_identifier_node.start_byte()..scoped_identifier_node.end_byte()];
                    return Some(scoped_identifier.to_string());
                }
            }

            // 一个文件只有一个packgeName
            return None;
        }

        if node_cursor.goto_first_child() {
            continue;
        }

        loop {
            if node_cursor.goto_next_sibling() {
                break;
            }

            if !node_cursor.goto_parent() {
                return None;
            }
        }
    }
}
